#ifndef PARSER_HPP
#define PARSER_HPP

#include <string>
#include <vector>
#include <memory>
#include <regex>
#include <map>
#include <iostream>
#include "instruction_types.hpp"

// Forward declarations
struct PTXMetadata;
struct PTXParameter;
struct PTXRegisterDeclaration;
struct PTXFunction;
struct PTXGlobalVariable;
struct PTXSymbolTable;
struct PTXProgram;

// ========================================
// PTX 元数据（文件级别）
// ========================================
struct PTXMetadata {
    std::string version;           // .version 6.0
    std::string target;            // .target sm_50
    int addressSize;               // .address_size 64
    bool debugMode;                // .target sm_50, debug
    
    PTXMetadata() : version(""), target(""), addressSize(64), debugMode(false) {}
};

// ========================================
// 参数定义
// ========================================
struct PTXParameter {
    std::string name;              // 参数名称
    std::string type;              // 参数类型：.u64, .s32, .f32
    size_t offset;                 // 在参数内存空间中的偏移量
    size_t size;                   // 参数大小（字节）
    bool isPointer;                // 是否是指针类型
    
    PTXParameter() : name(""), type(""), offset(0), size(0), isPointer(false) {}
    PTXParameter(const std::string& n, const std::string& t, size_t off, size_t sz, bool ptr = false)
        : name(n), type(t), offset(off), size(sz), isPointer(ptr) {}
};

// ========================================
// 寄存器声明
// ========================================
struct PTXRegisterDeclaration {
    std::string type;              // 寄存器类型：.s32, .u64, .f32, .pred
    std::string baseRegister;      // 基础寄存器名：%r, %f, %p, %rd
    int startIndex;                // 起始索引
    int count;                     // 数量
    
    PTXRegisterDeclaration() : type(""), baseRegister(""), startIndex(0), count(0) {}
    PTXRegisterDeclaration(const std::string& t, const std::string& base, int start, int cnt)
        : type(t), baseRegister(base), startIndex(start), count(cnt) {}
};

// ========================================
// 函数定义（包括 .func 和 .entry）
// ========================================
struct PTXFunction {
    std::string name;                          // 函数名称
    bool isEntry;                              // true = .entry (kernel), false = .func
    std::vector<PTXParameter> parameters;      // 输入参数列表
    std::vector<PTXParameter> returnValues;    // 返回值列表（仅 .func）
    std::vector<PTXRegisterDeclaration> registerDeclarations;
    size_t startInstructionIndex;              // 函数体第一条指令的索引
    size_t endInstructionIndex;                // 函数体最后一条指令的索引
    std::map<std::string, size_t> localLabels; // 标签名 → 指令索引
    
    PTXFunction() : name(""), isEntry(false), startInstructionIndex(0), endInstructionIndex(0) {}
};

// ========================================
// 全局变量/常量定义
// ========================================
struct PTXGlobalVariable {
    std::string name;              // 变量名称
    std::string type;              // 变量类型：.b8, .s32, .f32
    std::string space;             // 内存空间：.global, .shared, .const
    size_t alignment;              // 对齐要求
    size_t size;                   // 大小（字节）
    uint64_t address;              // 分配的内存地址
    std::vector<uint8_t> initialData;  // 初始化数据
    
    PTXGlobalVariable() : name(""), type(""), space(""), alignment(0), size(0), address(0) {}
};

// ========================================
// 符号表（全局）
// ========================================
struct PTXSymbolTable {
    std::map<std::string, PTXFunction> functions;
    std::map<std::string, size_t> globalLabels;
    std::map<std::string, PTXGlobalVariable> variables;
    std::map<std::string, PTXParameter*> parameterSymbols;
    
    PTXSymbolTable() {}
    
    const PTXFunction* findFunction(const std::string& name) const {
        auto it = functions.find(name);
        return (it != functions.end()) ? &(it->second) : nullptr;
    }
    
    bool findLabel(const std::string& name, const std::string& currentFunction, size_t& outIndex) const {
        auto funcIt = functions.find(currentFunction);
        if (funcIt != functions.end()) {
            auto labelIt = funcIt->second.localLabels.find(name);
            if (labelIt != funcIt->second.localLabels.end()) {
                outIndex = labelIt->second;
                return true;
            }
        }
        auto globalIt = globalLabels.find(name);
        if (globalIt != globalLabels.end()) {
            outIndex = globalIt->second;
            return true;
        }
        return false;
    }
    
    const PTXParameter* findParameter(const std::string& name) const {
        auto it = parameterSymbols.find(name);
        return (it != parameterSymbols.end()) ? it->second : nullptr;
    }
};

// ========================================
// 完整的 PTX 程序解析结果
// ========================================
struct PTXProgram {
    PTXMetadata metadata;
    std::vector<DecodedInstruction> instructions;
    PTXSymbolTable symbolTable;
    std::vector<PTXFunction> functions;
    std::vector<PTXGlobalVariable> globalVariables;
    std::vector<size_t> entryPoints;  // 索引到 functions 数组
    std::vector<std::string> errors;
    std::vector<std::string> warnings;
    
    PTXProgram() {}
    
    const PTXFunction* getMainEntry() const {
        if (entryPoints.empty()) return nullptr;
        return &functions[entryPoints[0]];
    }
    
    const PTXFunction* getEntryByName(const std::string& name) const {
        for (size_t idx : entryPoints) {
            if (functions[idx].name == name) {
                return &functions[idx];
            }
        }
        return nullptr;
    }
    
    void printSummary() const {
        std::cout << "=== PTX Program Summary ===" << std::endl;
        std::cout << "Version: " << metadata.version << std::endl;
        std::cout << "Target: " << metadata.target << std::endl;
        std::cout << "Address Size: " << metadata.addressSize << std::endl;
        std::cout << std::endl;
        
        std::cout << "Functions: " << functions.size() << std::endl;
        for (const auto& func : functions) {
            std::cout << "  " << (func.isEntry ? "[ENTRY]" : "[FUNC ]") 
                     << " " << func.name 
                     << " (" << func.parameters.size() << " params, "
                     << (func.endInstructionIndex - func.startInstructionIndex + 1) << " instructions)"
                     << std::endl;
        }
        std::cout << std::endl;
        
        std::cout << "Total Instructions: " << instructions.size() << std::endl;
        std::cout << "Global Variables: " << globalVariables.size() << std::endl;
        std::cout << "Entry Points: " << entryPoints.size() << std::endl;
        
        if (!errors.empty()) {
            std::cout << std::endl;
            std::cout << "Errors: " << errors.size() << std::endl;
            for (const auto& err : errors) {
                std::cout << "  ERROR: " << err << std::endl;
            }
        }
        
        if (!warnings.empty()) {
            std::cout << std::endl;
            std::cout << "Warnings: " << warnings.size() << std::endl;
            for (const auto& warn : warnings) {
                std::cout << "  WARNING: " << warn << std::endl;
            }
        }
    }
};

// ========================================
// PTXParser 类（重构后）
// ========================================
class PTXParser {
public:
    // Constructor/destructor
    PTXParser();
    ~PTXParser();

    // Parse a PTX file
    bool parseFile(const std::string& filename);

    // Parse PTX code from string
    bool parseString(const std::string& ptxCode);

    // Get parsed instructions (backward compatibility)
    const std::vector<DecodedInstruction>& getInstructions() const;

    // Get error message if parsing failed
    const std::string& getErrorMessage() const;
    
    // NEW: Get complete program structure
    const PTXProgram& getProgram() const;

private:
    // Private implementation details
    class Impl;
    std::unique_ptr<Impl> pImpl;
};

#endif // PARSER_HPP